//! Generic bitmap structure used to represent true/false for a chunk's blocks,
//! addressible using `SubchunkPos`.

use std::u32;

use craftdig_math::SubchunkPos;

use super::ChunkTable;

/// Size of the word we're representing the chunk bitmap with (always a u32).
pub const WORD_SIZE: u32 = u32::BITS;

/// Number of bits from a raw subchunk that we mask off.
pub const WORD_SUFFIX: usize = 5;

/// Mask from the raw subchunk that we use as index within a word.  Must correspond with `WORD_SUFFIX`.
pub const WORD_MASK: u16 = 0x001f;

/// Size of a chunk, but then broken down into
pub const WORDS: usize = (16 * 16 * 16) / WORD_SIZE as usize;

/// Bitmap addressible by `SubchunkPos` to represent a flag state for each block
/// in a chunk.
#[derive(Copy, Clone, Debug)]
pub struct ChunkBitmap {
    bits: [u32; WORDS],
}

impl ChunkBitmap {
    /// Creates a new instance with all bits unset.
    pub fn new_empty() -> Self {
        Self { bits: [0; WORDS] }
    }

    /// Creates a new instance with all bits set.
    pub fn new_full() -> Self {
        Self {
            bits: [u32::MAX; WORDS],
        }
    }

    /// Sets a bit true or false.
    pub fn set_bit_value(&mut self, pos: SubchunkPos, v: bool) {
        let (word_idx, subword_bit) = conv_scp_to_idxs(pos);
        let mut w = self.bits[word_idx];

        // Someone tell me if there's some fancy non-branching way to do this.
        if v {
            w |= 1 << subword_bit;
        } else {
            w &= !(1 << subword_bit)
        }

        self.bits[word_idx] = w;
    }

    /// Sets a bit to true.
    pub fn set(&mut self, pos: SubchunkPos) {
        self.set_bit_value(pos, true);
    }

    /// Sets a bit to false.
    pub fn unset(&mut self, pos: SubchunkPos) {
        self.set_bit_value(pos, false);
    }

    /// Checks a bit.
    pub fn check(&self, pos: SubchunkPos) -> bool {
        let (word_idx, subword_bit) = conv_scp_to_idxs(pos);
        let w = self.bits[word_idx];
        ((w >> subword_bit) & 1) == 1
    }

    /// Counts the total number of set bits.
    pub fn popcnt(&self) -> u32 {
        self.bits.iter().map(|wv| wv.count_ones()).sum()
    }

    /// Returns an iterator over the set bits in no particular order.
    pub fn iter_set(&self) -> impl Iterator<Item = SubchunkPos> + '_ {
        (0..WORDS).flat_map(|wi| {
            let mut wv = self.bits[wi];
            std::iter::from_fn(move || {
                if wv == 0 {
                    return None;
                } else {
                    let lz = wv.leading_zeros();
                    let leading_1_idx = WORD_SIZE - lz - 1;
                    wv = wv & !(1 << leading_1_idx);
                    Some(conv_idxs_to_scp(wi as u16, leading_1_idx as u16))
                }
            })
        })
    }

    // TODO iter_unset
}

impl Default for ChunkBitmap {
    fn default() -> Self {
        Self::new_empty()
    }
}

fn conv_scp_to_idxs(pos: SubchunkPos) -> (usize, u16) {
    let raw = pos.raw();
    let word_idx = raw >> WORD_SUFFIX;
    let subword_bit = raw & WORD_MASK;
    (word_idx as usize, subword_bit)
}

fn conv_idxs_to_scp(word_idx: u16, subword_bit: u16) -> SubchunkPos {
    assert!(
        (subword_bit & !WORD_MASK) == 0,
        "chunk_bitmap: invalid subword bit"
    );
    let raw = (word_idx << WORD_SUFFIX) | subword_bit;
    SubchunkPos::from_raw(raw)
}

impl ChunkTable for ChunkBitmap {
    type Value = bool;

    fn get_clone(&self, pos: SubchunkPos) -> Self::Value {
        self.check(pos)
    }

    fn set_swap(&mut self, pos: SubchunkPos, v: Self::Value) -> Self::Value {
        // As written this computes the coordinates twice, but LLVM is probably
        // smart enough to optimize that down.
        let vv = self.check(pos);
        self.set_bit_value(pos, v);
        vv
    }
}

#[cfg(test)]
mod tests {
    use craftdig_math::{SubchunkPos, CHUNK_NUM_BLOCKS};

    use super::ChunkBitmap;

    #[test]
    fn test_iter_full() {
        let bm = ChunkBitmap::new_full();
        assert_eq!(bm.popcnt(), CHUNK_NUM_BLOCKS as u32);
        assert_eq!(bm.iter_set().count(), CHUNK_NUM_BLOCKS);
    }

    #[test]
    fn test_iter_one() {
        let mut bm = ChunkBitmap::new_empty();
        let pos = SubchunkPos::new(11, 3, 8);
        bm.set(pos);
        let set = bm.iter_set().collect::<Vec<_>>();
        assert_eq!(set.as_slice(), &[pos]);
        assert_eq!(bm.popcnt(), 1);
    }

    #[test]
    fn test_iter_few() {
        let mut bm = ChunkBitmap::new_empty();
        let pos1 = SubchunkPos::new(1, 8, 3);
        let pos2 = SubchunkPos::new(5, 0, 0);
        let pos3 = SubchunkPos::new(9, 15, 4);
        bm.set(pos1);
        bm.set(pos2);
        bm.set(pos3);
        let set = bm.iter_set().collect::<Vec<_>>();
        assert_eq!(set.as_slice(), &[pos1, pos2, pos3]);
        assert_eq!(bm.popcnt(), 3);
    }
}
